from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.model_selection import train_test_split
import pandas as pd
import numpy as np

class DataPreprocessor:
    def __init__(self, test_size=0.3, random_state=42):
        self.test_size = test_size
        self.random_state = random_state
        self.scaler = StandardScaler()
        self.is_fitted = False
    
    def load_and_preprocess_iris(self):
        """加载并预处理鸢尾花数据集"""
        from sklearn.datasets import load_iris
        
        iris = load_iris()
        X = iris.data
        y = iris.target
        
        # 标准化特征
        X_scaled = self.scaler.fit_transform(X)
        self.is_fitted = True
        
        # 分割数据集
        X_train, X_test, y_train, y_test = train_test_split(
            X_scaled, y, 
            test_size=self.test_size, 
            random_state=self.random_state,
            stratify=y
        )
        
        return X_train, X_test, y_train, y_test, iris
    
    def preprocess_new_data(self, new_data):
        """预处理新数据（使用已拟合的scaler）"""
        if not self.is_fitted:
            raise ValueError("需要先调用load_and_preprocess_iris方法拟合数据")
        return self.scaler.transform(new_data)
